import sys
import os


# 获取项目根目录 (不这么搞的话，下面无法导包不了 
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '../src'))
sys.path.insert(0, project_root)
print(f'项目根目录 方便import测试: {project_root}')

from bt_config import BT_Config
from transformers import BertModel, BertConfig, AutoModel
from data_handle.data_loader import create_data_loader

params = BT_Config()
# 加载预训练BERT模型（可以替换成你本地的路径）
model = BertModel.from_pretrained(params.bert_path)
train_loader = create_data_loader(params.train_json_path, 2)

for step, batch in enumerate(train_loader):
    print('------------------')
    print(batch['input_ids'].shape)
    print(batch['attention_mask'].shape)
    print(batch['labels'].shape)
    
    output = model(batch['input_ids'], batch['attention_mask'])
    print('------------------')
    print(output)
    print('------------------')
    print(output.last_hidden_state.shape)
    print(output.pooler_output.shape)

    if step == 0 : break  # 先看第一批数据


# print(batch['input_ids'].shape)
# print(batch['attention_mask'].shape)
# print(batch['labels'].shape)
# torch.Size([2, 35])
# torch.Size([2, 35])
# torch.Size([2])
    

# print(output.last_hidden_state.shape)
# print(output.pooler_output.shape)
# torch.Size([2, 35, 768])
# torch.Size([2, 768])